package com.example.framework.web.core.filter;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.extra.servlet.JakartaServletUtil;
import com.example.framework.common.constant.BaseConstants;
import com.example.framework.web.core.util.TraceIdGenerateUtils;
import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.slf4j.MDC;
import org.springframework.web.filter.OncePerRequestFilter;

import java.io.IOException;

@Slf4j
public class TraceIdFilter extends OncePerRequestFilter {

    @Override
    protected void doFilterInternal(@NotNull HttpServletRequest request, @NotNull HttpServletResponse response, FilterChain chain) throws ServletException, IOException {
        String traceId = this.getTraceId(request);
        MDC.put(BaseConstants.TRACE_UNIQUE_ID, traceId);
        log.debug("http doFilter:[{}], uri:[{}]", this.getClass().getSimpleName(), request.getRequestURI());
        try {
            request.setAttribute(BaseConstants.TRACE_UNIQUE_ID, traceId);
            JakartaServletUtil.setHeader(response, BaseConstants.TRACE_UNIQUE_ID, traceId);
            chain.doFilter(request, response);
        } finally {
            MDC.clear();
        }
    }

    private String getTraceId(HttpServletRequest request) {
        Object atbTraceId = request.getAttribute(BaseConstants.TRACE_UNIQUE_ID);
        if (ObjectUtil.isNotEmpty(atbTraceId)) {
            return String.valueOf(atbTraceId);
        }
        String headerTraceId = request.getHeader(BaseConstants.TRACE_UNIQUE_ID);
        if (ObjectUtil.isNotEmpty(headerTraceId)) {
            return headerTraceId;
        }
        return TraceIdGenerateUtils.newTraceId();
    }

}
